'''
@author: zhangkai
@license: (C) Copyright 2017-2023
@contact: jeffcobile@gmail.com
@Software : PyCharm
@file: icnet.py
@time: 2020-06-09 11:03:07
@desc: 
'''
from ELib.model.model_zoo import MODEL_ZOO
from ELib.model.segbase import SegBaseModel
from ELib.basic_model import _ConvBNReLU
import torch


@MODEL_ZOO.register()
class ICNet(SegBaseModel):
    """Image Cascade Network"""

    def __init__(self, cfg):
        super(ICNet, self).__init__(cfg)
        self.conv_sub1 = torch.nn.Sequential(
            _ConvBNReLU(3, 32, 3, 2),
            _ConvBNReLU(32, 32, 3, 2),
            _ConvBNReLU(32, 64, 3, 2)
        )

        self.head = _ICHead(cfg, self.nclass)
        self.__setattr__('decoder', ['conv_sub1', 'head'])

    def forward(self, x):
        size = x.size()[2:]
        # sub 1
        x_sub1 = self.conv_sub1(x)

        # sub 2
        x_sub2 = torch.nn.functional.interpolate(x, scale_factor=0.5, mode='bilinear', align_corners=True)
        _, x_sub2, _, _ = self.encoder(x_sub2)

        # sub 4
        x_sub4 = torch.nn.functional.interpolate(x, scale_factor=0.25, mode='bilinear', align_corners=True)
        _, _, _, x_sub4 = self.encoder(x_sub4)

        outputs = self.head(x_sub1, x_sub2, x_sub4, size)

        return tuple(outputs)


class _ICHead(torch.nn.Module):
    def __init__(self, cfg, nclass, norm_layer=torch.nn.BatchNorm2d):
        super(_ICHead, self).__init__()
        scale = cfg.MODEL.BACKBONE_SCALE
        self.cff_12 = CascadeFeatureFusion(int(512 * scale), 64, 128, nclass, norm_layer)
        self.cff_24 = CascadeFeatureFusion(int(2048 * scale), int(512 * scale), 128, nclass, norm_layer)
        self.conv_cls = torch.nn.Conv2d(128, nclass, 1, bias=False)

    def forward(self, x_sub1, x_sub2, x_sub4, size):
        outputs = list()
        x_cff_24, x_24_cls = self.cff_24(x_sub4, x_sub2)
        outputs.append(x_24_cls)
        x_cff_12, x_12_cls = self.cff_12(x_sub2, x_sub1)
        outputs.append(x_12_cls)

        up_x2 = torch.nn.functional.interpolate(x_cff_12, scale_factor=2, mode='bilinear', align_corners=True)
        up_x2 = self.conv_cls(up_x2)
        outputs.append(up_x2)

        up_x8 = torch.nn.functional.interpolate(up_x2, size, mode='bilinear', align_corners=True)
        outputs.append(up_x8)
        # 1 -> 1/4 -> 1/8 -> 1/16
        outputs.reverse()

        return outputs


class CascadeFeatureFusion(torch.nn.Module):
    """CFF Unit"""

    def __init__(self, low_channels, high_channels, out_channels, nclass, norm_layer=torch.nn.BatchNorm2d):
        super(CascadeFeatureFusion, self).__init__()
        self.conv_low = torch.nn.Sequential(
            torch.nn.Conv2d(low_channels, out_channels, 3, padding=2, dilation=2, bias=False),
            norm_layer(out_channels)
        )
        self.conv_high = torch.nn.Sequential(
            torch.nn.Conv2d(high_channels, out_channels, 1, bias=False),
            norm_layer(out_channels)
        )
        self.conv_low_cls = torch.nn.Conv2d(out_channels, nclass, 1, bias=False)

    def forward(self, x_low, x_high):
        x_low = torch.nn.functional.interpolate(x_low, size=x_high.size()[2:], mode='bilinear', align_corners=True)
        x_low = self.conv_low(x_low)
        x_high = self.conv_high(x_high)
        x = x_low + x_high
        x = torch.nn.functional.relu(x, inplace=True)
        x_low_cls = self.conv_low_cls(x_low)

        return x, x_low_cls